'''
Created on Dec 10, 2009

@author: mkiyer
'''

import tables
import os
import numpy as np
import logging

from veggie.genome.chrom import get_genome
import veggie.io.wiggle as wiggle

class Track2(object):
    
    def __init__(self):
        self.group = None
        
class TrackDB(object):
    
    def __init__(self, filename):
        if os.path.exists(filename):
            if not tables.isHDF5File(filename):
                raise ValueError('file %s is not an HDF5 file' % filename)
            if not tables.isPyTablesFile(filename):
                raise ValueError('file %s is not a PyTables file' % filename)    
        self.h5file = tables.openFile(filename, 'r')

    def has_track(self, name):
        return ('/' + name) in self.h5file
        
    def get_track(self, name):
        if ('/' + name) in self.h5file:
            return self.h5file.getNode(self.h5file.root, name)


class TrackFactory(object):

    default_complevel = 1
    default_complib = 'lzo'
    default_filters = tables.Filters(complevel=default_complevel, 
                                     complib=default_complib,
                                     shuffle=True)
    
    # track types
    TRACK_TYPE_ARRAY = 0
    TRACK_TYPE_TABLE = 1 
    
    def __init__(self, filename):
        if os.path.exists(filename):
            if not tables.isHDF5File(filename):
                raise ValueError('file %s is not an HDF5 file' % filename)
            if not tables.isPyTablesFile(filename):
                raise ValueError('file %s is not a PyTables file' % filename)    
        self.h5file = tables.openFile(filename, 'a')

    def close(self):
        self.h5file.close()    

    def _create_track_arrays(self, group_name, dtype):
        logger = logging.getLogger(self.__class__.__name__)        
        # create subgroup
        if group_name is None:
            group = self.h5file.root
        elif ('/' + group_name) in self.h5file:
            logger.warning('h5file group %s already exists, data may be overwritten' % group_name)
            group = self.h5file.getNode('/' + group_name)
        else:
            group = self.h5file.createGroup(self.h5file.root, group_name)        
        # array datatype
        thisatom = tables.Atom.from_sctype(dtype)
        # create chromosome arrays
        for ref, length in get_genome().chrom_sizes.iteritems():
            ca = self.h5file.createCArray(group, ref,
                                          atom=thisatom,
                                          shape=(length,),
                                          filters=self.default_filters,
                                          chunkshape=None)
        # set attributes for group
        group._v_attrs.track_type = self.TRACK_TYPE_ARRAY
        return group

    def import_wiggle_as_array(self, wigglefile, track_name, dtype):        
        parent_group = self._create_track_arrays(track_name, dtype)
        wigglereader = wiggle.WiggleFileReader(open(wigglefile), dtype)
        for chunk in wigglereader.iterchunks():
            start, end, arr = chunk
            ca = self.h5file.getNode(parent_group, wigglereader.chrom)
            ca[start:end] = arr
    
    def import_tabular_as_array(self, bgfile, track_name, dtype,
                                chrom_col=1, start_col=2, end_col=3, value_col=4):
        logger = logging.getLogger(self.__class__.__name__)
        parent_group = self._create_track_arrays(track_name, dtype)
        # iterate through file
        debug_every = 1000000
        for linenum, line in enumerate(open(bgfile, 'r')):
            if linenum % debug_every == 0:
                logger.debug("%s: import_bedgraph_to_array %s finished %d" % (bgfile, track_name, linenum))        
            fields = line.strip().split('\t')
            chrom, start, end, pvalue = fields[chrom_col], int(fields[start_col]), int(fields[end_col]), int(fields[value_col])
            ca = self.h5file.getNode(parent_group, chrom)
            ca[start:end] = pvalue
    

class Track(object):
    
    def __init__(self, filename):
        if os.path.exists(filename):
            if not tables.isHDF5File(filename):
                raise ValueError('file %s is not an HDF5 file' % filename)
            if not tables.isPyTablesFile(filename):
                raise ValueError('file %s is not a PyTables file' % filename)    
        self.h5file = tables.openFile(filename, 'r')
    
    def close(self):
        self.h5file.close()
        
    def uniqueness(self, chrom, start, end):
        ca = self.h5file.getNode("/", chrom)
        return ca[start:end]

    @staticmethod
    def make_uniqueness_track(uniqueness_track, outfilename):
        h5file = tables.openFile(outfilename, 'w')
        default_complevel = 1
        default_complib = 'lzo'
        default_filters = tables.Filters(complevel=default_complevel, 
                                         complib=default_complib,
                                         shuffle=True)
        default_atom = tables.UInt8Atom()

        debug_every = 1000000
        for linenum, line in enumerate(open(uniqueness_track, 'r')):
            if linenum % debug_every == 0:
                print 'finished %d' % linenum        
            fields = line.strip().split('\t')
            chrom, start, end, pvalue = fields[1], int(fields[2]), int(fields[3]), int(fields[4])        
            if "/" + chrom not in h5file:
                length = get_genome().chrom_sizes[chrom]
                ca = h5file.createCArray(h5file.root, chrom,
                                         atom=default_atom,
                                         shape=(length,),
                                         filters=default_filters,
                                         chunkshape=None)
            else:
                ca = h5file.getNode("/", chrom)        
            ca[start:end] = pvalue
        h5file.close()


if __name__ == '__main__':
    import sys
    track_factory = TrackFactory(sys.argv[2])
    track_factory.import_wiggle_as_array(sys.argv[1], sys.argv[3], dtype=np.dtype(sys.argv[4]).type)    
    #Track.make_uniqueness_track(sys.argv[1], sys.argv[2])